refactor(speculative): one home each for the DFlash family's shared rules - #2757

### What does this PR do?

Type of change: new feature (performance), bug fix

Speeds up DFlash-family drafter training. On Gemma-4-E4B DSpark
streaming (PDX B300, 3 nodes = 1 vLLM serve + 2 trainers, 16 DP ranks,
batch 4 × 4096), step time goes **1.690 → 0.484 s (3.49×)**, with the
loss unchanged within run-to-run noise. Every round below was measured
on the same window (steps 21-40, profiler off).

| Change | Step time |
|---|---|
| FlexAttention over a block-sparse draft mask
(`dflash_use_flex_attention`, opt-in). The dense `[B, 1, Q, KV]` mask
had forced SDPA onto its sm80-only memory-efficient kernel. | 1.690 →
0.711 s |
| Offline/streaming and DSpark online: project only the base-teacher
rows a loss reads, not full `[B, seq, vocab]` logits; one device sync
for DSpark's metrics instead of five | 0.711 → 0.629 s |
| DSpark TVD chunked with `Tensor.split` instead of slicing, whose
backward zero-filled a full `[N, vocab]` tensor per chunk | 0.629 →
0.541 s |
| `torch.compile` of the draft stack and the DSpark TVD chunk
(`dflash_use_torch_compile`, opt-in). With compile or FlexAttention on,
the anchor count is fixed by the config and `seq_len`, so the compiled
shapes stay static: recompiles in 40 steps drop from 4 to 1. Eager runs
still size it to the batch. | 0.541 → 0.484 s |

The same rounds, by phase and by GPU kernel:

<img width="2040" height="2040" alt="image"
src="https://github.com/user-attachments/assets/13ae02e5-a171-4f48-8739-c6828d018413"
/>

Also:
- Offline/streaming LiLiCorr (#2589) reads its candidate target logits
the same per-row way, instead of rebuilding full-sequence logits.
- FlexAttention's conflicts with the attention sink and with dropout are
rejected at convert time.
- No-behavior-change cleanups: the attention visibility rule is written
once (`_draft_mask_mod`) for SDPA and FlexAttention, and KD is switched
by `dflash_self_logit_distillation` rather than by a nullable argument.

**Why not HF's `flex_attention` dispatch (the path EAGLE takes)?** The
BlockMask is built here either way; the difference is the kernel call.
HF's `flex_attention_forward` passes `enable_gqa=True` whenever the
query-head count is a power of two, and that backward took 568 ms on the
Gemma-4-E4B drafter: 10× the path that repeats K/V first, and 2.5×
slower than the SDPA baseline. At `head_dim` > 256, FlexAttention's
default tiles fail on sm103 (misaligned address, launch failure), so the
kernel tiles are pinned to 32×32 and the BlockMask block size has to
divide them; `dflash_flex_attention.py` keeps both rules together.

### Usage

```yaml
dflash:
  dflash_use_flex_attention: true
  dflash_use_torch_compile: true
```

On multi-node runs, point `TRITON_CACHE_DIR` and
`TORCHINDUCTOR_CACHE_DIR` at node-local paths: on a shared filesystem,
concurrent ranks race on the compile cache.

### Testing

- `tests/unit/torch/speculative`: 387 passed on transformers 5.12.
`tests/gpu/torch/speculative/plugins`: 14 passed, 1 skipped.
- FlexAttention on torch 2.8.0 (manual, RTX PRO 6000): against an fp64
dense reference on padded-row shapes (seq 128, Q_LEN 144 and 244),
forward and backward agree to ≤ 7e-7 relative in fp32, and
compute-sanitizer memcheck is clean. Before the `q_block` clamp,
memcheck caught the backward reading past `anchors`.
- New tests cover FlexAttention vs. dense (GPU), the mask rule against a
hand-written mask and on padded FlexAttention rows, TVD chunk-size
invariance and `split` in the backward graph, anchor sampling vs. the
legacy sampler in eager and static modes, the per-row teacher (lm_head
never sees the full sequence offline), and FlexAttention conflicts at
convert time and through a `from_pretrained` resume.
- The cleanups were checked old vs. new: losses, gradients, dense masks
and GPU BlockMasks are bitwise equal, except DSpark's rnn head, which
differs at bf16 rounding.
- Against main, offline DFlash KD, LiLiCorr and DSpark get bitwise-equal
losses and gradients in bf16 from the per-row teacher. The DFlash
regression recipe (eager) trains in the same time as on main with an
identical loss curve.
- Step times are medians over 16 ranks × steps 21-40 (profiler off) on
the setup above.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Both speedups are opt-in:
`dflash_use_flex_attention` and `dflash_use_torch_compile` default off.
The latter defaulted to True on main but did nothing; the shipped
recipes set it false.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌
- Did you get Claude approval on this PR?: ❌

### Additional Information

Based on main. The step times were measured on Gemma-4-E4B, which needs
#2186 to run; nothing in this PR is Gemma-specific.

The streaming `/done` fix that was part of this PR is now #2756.
Follow-up refactors are in draft #2757, stacked on this PR.

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added an optional FlexAttention path for DFlash draft attention, with
block-sparse masking.
* Added optional training-time compilation for draft model components;
compilation is disabled by default.
* Teacher logits can now be obtained from supplied logits or captured
hidden states during distillation.
* **Bug Fixes**
* Improved draft training behavior for batches with no valid anchors,
ensuring trainable parameters receive gradients.
* Added validation for FlexAttention configurations that do not support
attention sinks or dropout.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Hao Guo <haoguo@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…evice

DSpark's pseudo_speculative_generate adds its Markov bias, computed on
the draft's device, to the backbone logits that the base model's lm_head
returns. scripts/ar_validate.py loads with device_map="auto", and
accelerate keeps a tied lm_head with embed_tokens on the first GPU while
the draft can land on a later one. The logits then stay on the first
GPU, so the sum raises "Expected all tensors to be on the same device"
for every sample and AR validation reports nothing. This has been the
case since DSpark was added.

The backbone logits now move to the draft's device before the loop. On
a single device the move is a no-op.

Verified: a new GPU test dispatches the model with accelerate's
dispatch_model, embed_tokens and lm_head on the CPU and everything else
on cuda:0, and generates with DFlash and DSpark. DSpark fails without
the fix and passes with it. The bitwise A/B against #2757 shows DSpark's
drafted tokens unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

…nt and head

HFDominoModel did not override pseudo_speculative_generate, so AR
validation ran DFlash's, which decodes block positions 1..bs-1 with the
backbone alone. DFlash trains position k on the token at anchor + k;
Domino trains it on anchor + k + 1 (_label_shift = 1). Every drafted
token was therefore compared with the target one position too early,
so drafts were almost never accepted, and the correction head that
training optimizes and the exporter ships was never applied. domino.yaml
keeps AR validation off, so shipped training did not hit this.

DFlash's default decode now starts at block position 1 - _label_shift
and returns at most bs - 1 + _label_shift tokens. For DFlash, DFlash2
and LiLiCorr (shift 0) that is the same slice as before. Domino
overrides it the way DSpark does for its Markov head: it decodes the
block left to right, stepping the GRU over the anchor and the tokens
drafted so far, and from pure_draft_prefix_len on adds
embed_proj([hidden_k; GRU state]) to the backbone logits, as
_apply_domino_head does in training. As in DSpark's decode, the
backbone logits and the drafted tokens' embeddings are moved to the
draft's device: under device_map="auto" the embedding table, and a tied
lm_head with it, can sit on another GPU.

The eval forward warned that Domino's acceptance rates were
backbone-only, but that method computes no acceptance rate, and AR
validation now runs the head. The warning goes, and so do the other
copies of the claim: the comments in domino.yaml, dspark.yaml and the
Nemotron-3.5 DSpark warm-start recipe, the note in the Domino launcher
example, and the speculative-decoding skill's Domino, DSpark and
validation sheets. DSpark's generation has applied its Markov head
since DSpark was added. The validation sheet keeps its rule against
reporting such an AR as the model's result, now for the heads that AR
validation still skips (LiLiCorr's reranker, DFlash2's selector), and
adds the LiLiCorr and online DFlash2 examples to the ones it lists as
shipping no benchmark task. The Domino sheet and example now say how to
add the HF AR eval: DFlash's task_2 with --steps equal to the block
size, since Domino drafts the whole block. The flag values do not
change.

Verified: a new test checks that Domino's greedy drafts equal the
argmax of _apply_domino_head teacher-forced on them, with head weights
large enough to overturn the backbone. It fails on the parent commit,
and also when the head is skipped, its GRU state lags a step, or the
prefix boundary is off by one. The GPU test that puts lm_head on
another device now covers Domino as well, and fails without the move.
Bitwise A/B against #2757: DFlash and DSpark loss, gradients and drafted
tokens are unchanged; only Domino's drafted tokens change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>